#!/bin/bash
source "${LMD_BASE_INSTALL_SCRIPT_DIR}/global/sys_check.sh"

ensure_package() {
    local package_name=$1
    local package_version=$2

    # 特殊处理 insightface：使用指定的 whl 包安装
    if [ "$package_name" = "insightface" ]; then
        if conda_run_cmd pip show "$package_name" &> /dev/null; then
            echo "$package_name is already installed"
            return 0
        fi
        
        echo "$package_name is not installed. Installing with custom whl..."
        conda_run_pip_install "${GIT_HOST}/skippzz/insightface-0.7.3-cp312-cp312-win_amd64/blob/main/insightface-0.7.3-cp312-cp312-win_amd64.whl"
        
        if conda_run_cmd pip show "$package_name" &> /dev/null; then
            echo "[INFO] Successfully installed $package_name (custom whl)"
            return 0
        else
            echo "[ERROR] Failed to install $package_name (custom whl)"
            return 1
        fi
    fi

    # 其他包的默认安装逻辑
    if conda_run_cmd pip show "$package_name" &> /dev/null; then
        echo "$package_name is already installed"
    else
        echo "$package_name is not installed. Installing now..."

        if [ -n "$package_version" ]; then
            conda_run_pip_install "$package_name==$package_version"
        else
            conda_run_pip_install "$package_name"
        fi

        if conda_run_cmd pip show "$package_name" &> /dev/null; then
            echo "[INFO] Successfully installed $package_name"
            return 0
        else
            echo "[ERROR] Failed to install $package_name"
            return 1
        fi
    fi
}

check_nunchaku() {
    if current_os_is_windows; then
        if command -v nvidia-smi &> /dev/null; then
            # check nunchaku dir?
            # NUNCHAKU_DIR="$CURRENT_APP_PATH/models/ComfyUI-nunchaku"
            # echo "$NUNCHAKU_DIR exists."
            ensure_package insightface
            ensure_package xformers 0.0.31
            ensure_package sageattention 1.0.6
            ensure_package facexlib
            
            if conda_run_cmd pip show nunchaku &> /dev/null; then
                echo "nunchaku is installed"
            else
                echo "start install nunchaku"
                conda_run_pip_install ${GIT_HOST}/nunchaku-tech/nunchaku/releases/download/v1.0.1/nunchaku-1.0.1+torch2.7-cp312-cp312-win_amd64.whl
            fi
        fi
    fi
}

check_nunchaku